from typing import Any

import torch

import torchair
from torchair import register_fx_node_ge_converter
from torchair.ge import Tensor


@register_fx_node_ge_converter(torch.ops.npu.npu_grouped_matmul_swiglu_quant.default)
def convert_npu_grouped_matmul_swiglu_quant(
    x: Tensor,
    weight: Tensor,
    group_list: Tensor,
    weight_scale: Tensor,
    x_scale: Tensor,
    output: Tensor = None,
    output_scale: Tensor = None,
    output_offset: Tensor = None,
    meta_outputs: Any = None
):
    return torchair.ge.custom_op(
        "GroupedMatmulSwigluQuant",
        inputs={
            "x": x,
            "weight": weight,
            "weight_scale": weight_scale,
            "x_scale": x_scale,
            "group_list": group_list
        },
        attrs={
        },
        outputs=['output', 'output_scale', 'output_offset']
    )